import numpy as np
from scipy.integrate import solve_ivp
from math import sin, cos, radians
import matplotlib.pyplot as plt
# 引入 Matplotlib 的字体管理器
from matplotlib import rcParams

# --- 设置中文显示 ---
# 尝试使用系统自带的常见中文字体
# 如果系统没有这些字体，可能需要安装或更改为其他字体名称
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Heiti TC'] 
plt.rcParams['axes.unicode_minus'] = False # 解决负号 '-' 显示为方块的问题
# --------------------

# --- I. 物理常数和弹丸参数 ---
# 国际标准大气模型参数
RHO0 = 1.225
T0 = 288.15
L = 0.0065
G = 9.80
R = 287.05
P_TERM = (G / (R * L)) - 1 

# 弹丸参数
C_SOUND = 338.0
M = 5.0
A = 0.0095

# --- II. 目标参数 ---
TARGET_RANGE = 1200.0
TEST_ANGLE = 45.0
V0_GIVEN = 145.020 # 指定的初速度 (m/s)

# --- III. Cd-Ma 和 Rho-z 关系函数 ---

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5, 10.0]) 
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40, 0.40])

def get_drag_coefficient(v, c_sound):
    """根据马赫数线性插值计算 Cd。"""
    mach_number = v / c_sound
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS)
    return Cd

def get_air_density(z):
    """根据国际标准大气模型计算高度 z 处的空气密度 rho(z)。"""
    h = max(0.0, z)
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0:
        return 0.0
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

# --- IV. 弹道微分方程 (动态 ρ(z) 和 Cd(Ma)) ---

def equations_of_motion(t, S, theta, v0, params):
    """弹道微分方程组。S = [x, z, vx, vz]"""
    x, z, vx, vz = S
    m, A, g, c_sound = params

    v = np.sqrt(vx**2 + vz**2)
    rho = get_air_density(z)
    Cd = get_drag_coefficient(v, c_sound)
    
    drag_coeff = -(Cd * rho * A) / (2.0 * m) 
    
    ax = drag_coeff * v * vx
    az = drag_coeff * v * vz - g
    
    return [vx, vz, ax, az]

# --- V. 轨迹积分函数 ---

def calculate_trajectory(v0, theta):
    """
    给定初速 v0 和角度 theta，计算并返回完整的轨迹数据。
    """
    params = (M, A, G, C_SOUND)
    theta_rad = radians(theta)
    S0 = [
        0.0, 0.0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ]
    
    def hit_ground(t, S, theta, v0, params):
        return S[1]
    
    hit_ground.terminal = True 
    hit_ground.direction = -1 
    
    t_span = [0, 500] 
    
    sol = solve_ivp(
        equations_of_motion, 
        t_span, 
        S0, 
        args=(theta, v0, params), 
        method='RK45', 
        events=hit_ground,
        rtol=1e-7, 
        atol=1e-9,
        dense_output=True
    )
    
    if len(sol.y_events[0]) > 0:
        final_x = sol.y_events[0][0][0]
        t_hit = sol.t_events[0][0]
        t_plot = np.linspace(0, t_hit, 500)
        z_plot = sol.sol(t_plot)
        
        return final_x, z_plot[0], z_plot[1] 
    else:
        return -1.0, None, None

# --- VI. 执行计算和绘图 ---

print(f"--- 🚀 正在使用初速度 v₀ = {V0_GIVEN:.3f} m/s 计算轨迹 ---")

# 1. 计算完整轨迹
R_final, x_traj, z_traj = calculate_trajectory(V0_GIVEN, TEST_ANGLE)

if R_final < 0:
    print("❌ 轨迹计算失败，弹丸未落地或积分错误。")
else:
    # 2. 提取关键信息
    max_height = np.max(z_traj)
    max_height_idx = np.argmax(z_traj)

    print("\n--- 🏁 轨迹计算结果总结 ---")
    print(f"设定初速度 v₀: {V0_GIVEN:.3f} m/s")
    print(f"计算得到的射程 R: {R_final:.3f} m")
    print(f"最大射高 Z_max: {max_height:.3f} m")

    # 3. 绘制轨迹
    plt.figure(figsize=(10, 6))
    
    # 绘图标签和标题使用中文
    plt.plot(x_traj, z_traj, 
             label=f'初速度 $v_0 = {V0_GIVEN:.3f} \,\, \mathrm{{m/s}}$', 
             color='blue', linewidth=2)
    
    # 绘制目标点和地面线
    plt.plot(TARGET_RANGE, 0, 'ro', markersize=8, 
             label=f'目标点 ({TARGET_RANGE:.0f} m)')
    plt.hlines(0, 0, R_final * 1.05, color='gray', linestyle='--', linewidth=0.8, label='地面')
    
    # 标记落点
    plt.plot(R_final, 0, 'go', markersize=8, 
             label=f'实际落点 ({R_final:.1f} m)')

    # 标记最大射高
    plt.plot(x_traj[max_height_idx], max_height, 'kx', markersize=8, mew=2, 
             label=f'最大射高 ({max_height:.1f} m)')
    
    # 设置中文标题和坐标轴标签
    plt.title(f'弹丸飞行轨迹 ($\theta={TEST_ANGLE}^\circ$, 动态阻力/密度)', fontsize=14)
    plt.xlabel('水平距离 x (m)', fontsize=12)
    plt.ylabel('高度 z (m)', fontsize=12)
    plt.grid(True, linestyle=':', alpha=0.6)
    plt.legend()
    plt.gca().set_aspect('equal', adjustable='box') 
    plt.ylim(bottom=0)
    plt.xlim(right=max(TARGET_RANGE, R_final) * 1.1)
    
    plt.show()

print("\n--- ❗ 绘图总结 ❗ ---")
print("已在程序中引入 Matplotlib 中文设置。如果本地环境配置了'SimHei'等字体，中文标题和标签将正常显示。")
print(f"使用初速度 v₀ = {V0_GIVEN:.3f} m/s 计算得到的射程为 {R_final:.3f} m。")